#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MD文件数据清洗脚本
功能：保留文章高信息比率的内容，去掉摘要之前、参考文献之后的所有内容
支持：中文和英文文献的自动识别和处理
"""

import re
import os
import sys

def clean_md_file(file_path):
    """
    清洗MD文件，去掉摘要及之前的内容、参考文献及之后的内容，只保留正文
    支持中文和英文文献的自动识别和处理
    
    Args:
        file_path (str): MD文件路径
    
    Returns:
        str: 清洗后的内容
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"读取文件失败: {e}")
        return None
    
    lines = content.split('\n')
    
    # 直接查找绪论/引言作为正文开始位置
    content_start = -1
    
    # 定义绪论/引言的各种可能表达形式（只匹配以#开头的标题格式）
    intro_patterns = [
        # 中文模式
        r'^#\s*[0-9]*\s*绪\s*论',  # # 绪论、# 1 绪论等
        r'^#\s*[0-9]*\s*引\s*言',  # # 引言、# 1 引言等
        r'^#\s*[0-9]*\s*前\s*言',  # # 前言
        r'^#\s*[0-9]*\s*概\s*述',  # # 概述
        r'^#\s*[0-9]*\s*导\s*论',  # # 导论
        r'^#\s*[0-9]*\s*序\s*言',  # # 序言
        r'^#\s*第[一二三四五六七八九十0-9]+章\s+(绪论|引言|前言|概述|导论|序言)',  # # 第一章 绪论等
        r'^#\s*第\s*[一二三四五六七八九十0-9]+\s*章\s*绪\s*论',  # # 第 1 章 绪论等
        # 英文模式
        r'^#\s*[0-9]*\s*Introduction',  # # Introduction、# 1 Introduction等
        r'^#\s*[0-9]*\s*Preface',  # # Preface
        r'^#\s*[0-9]*\s*Overview',  # # Overview
        r'^#\s*[0-9]*\s*Background',  # # Background
        r'^#\s*[0-9]*\s*Motivation',  # # Motivation
        r'^#\s*[0-9]*\s*Prologue',  # # Prologue
        r'^#\s*[0-9]*\s*Foreword',  # # Foreword
        r'^#\s*Chapter\s*[0-9]+\s*:?\s*(Introduction|Background|Overview)',  # # Chapter 1: Introduction等
        r'^#\s*[0-9]+\.\s*(Introduction|Background|Overview)',  # # 1. Introduction等
    ]
    
    # 首先查找目录的结束位置
    toc_end = -1
    toc_start = -1
    for i, line in enumerate(lines):
        line_stripped = line.strip()
        # 查找目录标题（中英文）
        if re.search(r'^#*\s*目\s*录', line_stripped) or \
           re.search(r'^#*\s*Contents?', line_stripped, re.IGNORECASE) or \
           re.search(r'^#*\s*Table\s+of\s+Contents', line_stripped, re.IGNORECASE):
            print(f"找到目录标题在第{i+1}行: {line_stripped}")
            toc_start = i
            # 从目录开始位置向后查找，直到找到目录结束
            for j in range(i + 1, len(lines)):
                next_line = lines[j].strip()
                # 如果遇到明显的非目录内容（如图表目录、符号说明等），认为目录结束
                if re.search(r'^#*\s*图\s*目\s*录', next_line) or \
                   re.search(r'^#*\s*表\s*目\s*录', next_line) or \
                   re.search(r'^#*\s*符号', next_line) or \
                   re.search(r'^#*\s*List\s+of\s+(Figures?|Tables?|Symbols?)', next_line, re.IGNORECASE) or \
                   re.search(r'^#*\s*Symbols?', next_line, re.IGNORECASE) or \
                   re.search(r'^#*\s*Nomenclature', next_line, re.IGNORECASE) or \
                   re.search(r'^#*\s*Abbreviations?', next_line, re.IGNORECASE):
                    toc_end = j
                    print(f"目录结束于第{j+1}行: {next_line}")
                    break
                # 如果遇到章节标题（如# 1 绪论），认为目录结束
                elif re.search(r'^#\s*[0-9]+', next_line):
                    toc_end = j - 1  # 目录结束在章节标题之前
                    print(f"目录结束于第{j}行（遇到章节标题）: {next_line}")
                    break
                # 如果遇到绪论/引言等关键词，也认为目录结束
                elif any(re.search(pattern, next_line, re.IGNORECASE) for pattern in intro_patterns):
                    toc_end = j - 1  # 目录结束在绪论/引言之前
                    print(f"目录结束于第{j}行（遇到绪论/引言）: {next_line}")
                    break
            break
    
    print(f"目录结束位置: {toc_end}")
    
    # 查找绪论/引言开始位置（跳过目录部分）
    content_start = -1
    search_start = max(0, toc_end) if toc_end != -1 else 0
    print(f"开始从第{search_start+1}行搜索正文")
    
    for i in range(search_start, len(lines)):
        line_stripped = lines[i].strip()
        # 检查是否匹配任何绪论/引言模式
        for pattern in intro_patterns:
            if re.search(pattern, line_stripped, re.IGNORECASE):
                print(f"在第{i+1}行找到可能的绪论/引言: {line_stripped}")
                # 确保这不是目录中的条目，检查是否在目录部分
                # 如果当前位置在目录结束之后，或者没有找到目录，则认为是正文开始
                if toc_end == -1 or i > toc_end:
                    print(f"位置在目录之后，确认为正文开始")
                    content_start = i
                    print(f"确定正文开始位置: 第{i+1}行")
                    break
                else:
                    print(f"  位置在目录内，继续搜索")
        if content_start != -1:
            break
    
    # 如果找不到绪论/引言，查找第一个章节标题
    if content_start == -1:
        print("未找到绪论/引言，查找第一个章节标题")
        # 先跳过摘要部分，查找第一个章节标题
        abstract_end = -1
        for i, line in enumerate(lines):
            line_stripped = line.strip()
            # 查找摘要结束或关键词结束的位置（中英文）
            if re.search(r'^(Key\s*words?|Keywords?)\s*[:：]', line_stripped, re.IGNORECASE) or \
               re.search(r'^(关键词)\s*[:：]', line_stripped) or \
               re.search(r'^Abstract\s*$', line_stripped, re.IGNORECASE) or \
               re.search(r'^摘\s*要\s*$', line_stripped):
                # 找到关键词行，继续向下找到摘要结束
                for j in range(i + 1, min(i + 10, len(lines))):
                    next_line = lines[j].strip()
                    if next_line == '' or re.search(r'^#\s*[0-9]', next_line) or \
                       re.search(r'^[0-9]+[\s\.。]', next_line):
                        abstract_end = j
                        break
                if abstract_end != -1:
                    break
        
        # 从摘要结束位置之后查找第一个章节标题，但要确保不在目录中
        search_start = max(0, abstract_end) if abstract_end != -1 else 0
        # 如果有目录，确保搜索从目录结束之后开始
        if toc_end != -1:
            search_start = max(search_start, toc_end + 1)
        
        for i in range(search_start, len(lines)):
            line_stripped = lines[i].strip()
            # 查找第一个章节标题（中英文格式），但必须是以#开头的标题格式
            if re.search(r'^#\s*[0-9]', line_stripped) or \
               re.search(r'^#\s*第[0-9一二三四五六七八九十]+章', line_stripped) or \
               re.search(r'^#\s*[一二三四五六七八九十]\s*[、.]', line_stripped) or \
               re.search(r'^#\s*Chapter\s*[0-9]+', line_stripped, re.IGNORECASE) or \
               re.search(r'^#\s*[0-9]+\.', line_stripped):
                content_start = i
                print(f"找到第一个章节标题: 第{i+1}行 - {line_stripped}")
                break
            # 如果是纯数字开头的行（如目录条目），跳过
            elif re.search(r'^[0-9]+[\s\.。]', line_stripped) and not line_stripped.startswith('#'):
                print(f"跳过目录条目: 第{i+1}行 - {line_stripped}")
                continue
    
    # 如果还是找不到，从文件开头开始
    if content_start == -1:
        print("警告：未找到明确的正文开始位置，从文件开头开始保留")
        content_start = 0
    
    # 查找参考文献开始位置（只在正文开始位置之后搜索）
    reference_start = len(lines)  # 默认到文件末尾
    search_start_for_ref = content_start if content_start != -1 else 0
    for i in range(search_start_for_ref, len(lines)):
        line_stripped = lines[i].strip()
        # 匹配参考文献标题（包含各种可能的格式，中英文）
        if re.search(r'^#\s*[0-9]*\s*参\s*考\s*文\s*献', line_stripped) or \
           re.search(r'^#\s*[0-9]*\s*References?', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Bibliography', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Works?\s+Cited', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Literature\s+Cited', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*ORCID', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Appendix', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Supplementary', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Supporting\s+information', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*\[\s*参\s*考\s*文\s*献\s*\]', line_stripped) or \
           re.search(r'^参\s*考\s*文\s*献', line_stripped) or \
           re.search(r'^References?', line_stripped, re.IGNORECASE) or \
           re.search(r'^Bibliography', line_stripped, re.IGNORECASE) or \
           re.search(r'^Literature\s+Cited', line_stripped, re.IGNORECASE) or \
           re.search(r'^ORCID', line_stripped, re.IGNORECASE) or \
           re.search(r'^Appendix', line_stripped, re.IGNORECASE) or \
           re.search(r'^Supplementary', line_stripped, re.IGNORECASE) or \
           re.search(r'^Supporting\s+information', line_stripped, re.IGNORECASE) or \
           re.search(r'^\[\s*[0-9]+\s*\]', line_stripped):
            reference_start = i  # 不包含参考文献标题行
            print(f"找到参考文献开始位置: 第{i+1}行 - {line_stripped}，正文将在此行之前结束")
            break
    
    # 查找附录/附件开始位置（只在正文开始位置之后搜索）
    appendix_start = len(lines)  # 默认到文件末尾
    for i in range(search_start_for_ref, len(lines)):
        line_stripped = lines[i].strip()
        # 匹配附录、附件标题（包含各种可能的变体，中英文）
        if re.search(r'^#\s*[0-9]*\s*附\s*[录件表图]', line_stripped) or \
           re.search(r'^#\s*[0-9]*\s*附\s*录\s*件', line_stripped) or \
           re.search(r'^#\s*[0-9]*\s*Append(ix|ices?)\s*[A-Z]?\s*$', line_stripped, re.IGNORECASE) or \
           re.search(r'^#\s*[0-9]*\s*Supplement(ary)?\s*(Material|Data)?', line_stripped, re.IGNORECASE) or \
           re.search(r'^附\s*[录件表图]\s*$', line_stripped) or \
           re.search(r'^附\s*录\s*件\s*$', line_stripped) or \
           re.search(r'^Append(ix|ices?)\s*[A-Z]?\s*$', line_stripped, re.IGNORECASE):
            appendix_start = i
            print(f"找到附录开始位置: 第{i+1}行 - {line_stripped}")
            break
    
    # 取最早出现的结束位置
    content_end = min(reference_start, appendix_start)
    print(f"正文开始位置: {content_start+1}, 正文结束位置: {content_end+1}")
    
    # 提取正文内容
    if content_start >= content_end:
        print(f"警告：未找到有效的正文内容 (开始位置 {content_start+1} >= 结束位置 {content_end+1})")
        return ""
    
    cleaned_lines = lines[content_start:content_end]
    
    # 去除空行过多的情况，保持适当的段落间距
    result_lines = []
    empty_line_count = 0
    
    for line in cleaned_lines:
        if line.strip() == '':
            empty_line_count += 1
            if empty_line_count <= 2:  # 最多保留2个连续空行
                result_lines.append(line)
        else:
            empty_line_count = 0
            result_lines.append(line)
    
    return '\n'.join(result_lines)

def process_single_file(input_file, output_dir=None):
    """
    处理单个文件
    
    Args:
        input_file (str): 输入文件路径
        output_dir (str): 输出目录，如果为None则使用默认cleaned文件夹
    """
    if not os.path.exists(input_file):
        print(f"错误：文件不存在 {input_file}")
        return False
    
    cleaned_content = clean_md_file(input_file)
    if cleaned_content is None:
        return False
    
    # 确定输出文件路径
    if output_dir is None:
        output_dir = os.path.join(os.path.dirname(input_file), "cleaned")
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取文件名
    filename = os.path.basename(input_file)
    output_file = os.path.join(output_dir, filename)
    
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(cleaned_content)
        print(f"成功处理文件: {input_file} -> {output_file}")
        return True
    except Exception as e:
        print(f"写入文件失败: {e}")
        return False

def process_directory(directory_path, output_dir=None):
    """
    批量处理目录中的所有MD文件
    
    Args:
        directory_path (str): 目录路径
        output_dir (str): 输出目录，如果为None则使用默认cleaned文件夹
    """
    if not os.path.exists(directory_path):
        print(f"错误：目录不存在 {directory_path}")
        return
    
    md_files = [f for f in os.listdir(directory_path) if f.endswith('.md')]
    
    if not md_files:
        print("目录中没有找到MD文件")
        return
    
    print(f"找到 {len(md_files)} 个MD文件")
    
    # 确定输出目录
    if output_dir is None:
        output_dir = os.path.join(directory_path, "cleaned")
    
    print(f"输出目录: {output_dir}")
    
    success_count = 0
    for md_file in md_files:
        file_path = os.path.join(directory_path, md_file)
        if process_single_file(file_path, output_dir):
            success_count += 1
    
    print(f"\n处理完成！成功处理 {success_count}/{len(md_files)} 个文件")

def main():
    """
    主函数
    """
    print("MD文件数据清洗工具")
    print("=" * 50)
    
    if len(sys.argv) > 1:
        # 命令行模式
        target_path = sys.argv[1]
        output_dir = sys.argv[2] if len(sys.argv) > 2 else None
        
        if os.path.isfile(target_path):
            # 处理单个文件
            process_single_file(target_path, output_dir)
        elif os.path.isdir(target_path):
            # 处理目录
            process_directory(target_path, output_dir)
        else:
            print(f"错误：路径不存在 {target_path}")
    else:
        # 交互模式
        print("请选择处理模式：")
        print("1. 处理单个文件")
        print("2. 批量处理目录中的所有MD文件")
        
        choice = input("请输入选择 (1/2): ").strip()
        
        if choice == '1':
            file_path = input("请输入MD文件路径: ").strip()
            output_dir = input("请输入输出目录路径 (回车使用默认cleaned文件夹): ").strip()
            output_dir = output_dir if output_dir else None
            process_single_file(file_path, output_dir)
            
        elif choice == '2':
            dir_path = input("请输入目录路径: ").strip()
            output_dir = input("请输入输出目录路径 (回车使用默认cleaned文件夹): ").strip()
            output_dir = output_dir if output_dir else None
            process_directory(dir_path, output_dir)
            
        else:
            print("无效选择")

if __name__ == "__main__":
    main()